# model_predictor.py
import os

import joblib
from config import MODEL_SAVE_PATH, OUTPUT_FILE

class ModelPredictor:
    def __init__(self):
        self.model = None

    def load_model(self):
        if not self.model:
            if os.path.exists(MODEL_SAVE_PATH):
                self.model = joblib.load(MODEL_SAVE_PATH)
            else:
                raise FileNotFoundError("未找到训练好的模型，请先训练模型")
        return self.model

    def predict(self, texts):
        model = self.load_model()
        predictions = model.predict(texts)

        with open(OUTPUT_FILE, 'w', encoding='utf-8') as f:
            for pred in predictions:
                f.write(f"{pred}\n")

        print(f"预测结果已保存到: {OUTPUT_FILE}")
        return predictions